function [slab_samples,scale_delta_samples,llik_samples] = sample_beta_reg(y,X,burnin,ngibbs)
%
% [slab,scale_delta,llik] = sample_beta_reg(y,X,burnin,ngibbs)
%
%  y           = n x 1 response vector (0, 1)
%  X           = n x p predictor matrix
%  scale       = 1 x 1 scale of beta distribution
%
%  slab        = p x ngibbs
%  scale_delta = p x ngibbs
%  llik        = 1 x (ngibbs + burnin)
%
% code: Yongjin Park, ypp@csail.mit.edu
%
    
    [n,~] = size(X);

    % add intercept
    X = [X, ones(n,1)];

    [~,p] = size(X);

    assert(size(y,1) == n && size(y,2) == 1, 'check if y = n x 1');

    % check extreme value y = 0 or y = 1
    extreme = 1e-4;
    if any(y(:) < extreme) || any(y(:) > 1-extreme),
        y = (y - 0.5)*(1-extreme) + 0.5;
    end


    scale_delta_samples = zeros(p,ngibbs,'single');
    slab_samples = zeros(p,ngibbs,'single');
    llik_samples = zeros(1,ngibbs+burnin,'single');

    slab = zeros(p,1,'single');

    nslice = 1;
    slice_width = 0.5;

    TOL = 1e-4;

    SLAB_LOWER = -5;
    SLAB_UPPER = 5;

    scale_delta = zeros(p,1);
    scale = exp(X*scale_delta); % scale ~ exp(1+delta) and delta ~ N(0,1)
    scale_delta_samples(p,ngibbs);

    for iter = 1:(burnin + ngibbs),

        % ================================================================
        % 1. Sample slab by slice & gibbs sampling

        % evaluated log density at current slab
        for k = 1:p,
            fval = llik_func(y,X,slab,scale);
            fval_prime = fval + log(rand);

            % (a) step-out
            left = slab;
            right = slab;
            slab_prime = slab;

            rr = rand;
            left(k) = left(k) - rr * slice_width;
            right(k) = right(k) + (1-rr) * slice_width;

            while (llik_func(y,X,left,scale) > fval_prime) && left(k) > SLAB_LOWER,
                left(k) = left(k) - slice_width;
            end
            left(k) = max(SLAB_LOWER, left(k));

            while (llik_func(y,X,right,scale) > fval_prime) && right(k) < SLAB_UPPER,
                right(k) = right(k) + slice_width;
            end
            right(k) = min(SLAB_UPPER, right(k));

            % (b) shrinkage
            while 1
                slab_prime(k) = rand() * (right(k) - left(k)) + left(k);
                fval = llik_func(y,X,slab_prime,scale);

                if fval >= fval_prime - TOL,
                    break;
                else
                    if slab_prime(k) > slab(k)
                        right(k) = slab_prime(k);
                    elseif slab_prime(k) < slab(k)
                        left(k) = slab_prime(k);
                    else
                        error('error in slab sampling');
                    end
                end
            end
            slab = slab_prime;


        end


        % ================================================================
        % 2. Sample scale_delta model
        for k = 1:p,
            fval = llik_func(y,X,slab,scale) - 0.5*sum(scale_delta.^2);
            fval_prime = fval + log(rand);

            % step out
            left = scale_delta;
            right = scale_delta;
            scale_delta_prime = scale_delta;

            rr = rand;
            left(k) = left(k) - rr * slice_width;
            right(k) = right(k) + (1-rr) * slice_width;

            % pick a point by shrinking
            while 1
                scale_delta_prime(k) = rand() * (right(k) - left(k)) + left(k);
                scale_prime = exp(X*scale_delta_prime);
                
                fval = llik_func(y,X,slab,scale_prime) - 0.5*sum(scale_delta_prime.^2);
                
                if fval >= fval_prime - TOL,
                    break;
                else
                    if scale_delta_prime(k) > scale_delta(k)
                        right(k) = scale_delta_prime(k);
                    elseif scale_delta_prime(k) < scale_delta(k)
                        left(k) = scale_delta_prime(k);
                    else
                        error('error in scale sampling');
                    end
                end
            end
            
            scale_delta = scale_delta_prime;
            scale = scale_prime;
        end
        

        % keep track of samples
        if iter > burnin,
            slab_samples(:,iter - burnin) = slab(1:p);
            scale_delta_samples(:,iter-burnin) = scale_delta;
        end

        llik = llik_func(y,X,slab_prime,scale);
        llik_samples(iter) = llik;
        fprintf('iter = %05d, llik = %.4e\r',iter,llik);

    end

end

function fval = llik_func(y,X,slab,scale)

    TOL = 1e-4;

    assert(all(scale > 0), 'must be strictly positive scale');

    mu = 1./(1+exp(- X * slab));
    mu(mu < TOL) = TOL;
    mu(mu > 1 - TOL) = 1 - TOL;
    a = mu.*scale;
    b = (1-mu).*scale;

    fval = sum(gammaln(scale) -gammaln(a) - gammaln(b) + (a-1).*log(y) + (b-1).*log(1-y) );
end
